# -*- coding: utf-8 -*-

import time
import math
from bson import ObjectId
import arrow
from pymongo import MongoClient
from conf.settings import CLOUDMONITOR_MONGODB, CLOUDMONITOR_DB, WASS_DB
from utility.log import Logger
from utility.monitor import Monitor
from utility.wx_notify import notify
monitor = Monitor()
server = MongoClient(CLOUDMONITOR_MONGODB)
cloud = server[CLOUDMONITOR_DB]
wass = server[WASS_DB]
logger = Logger('check_task.log').get_logger()
notify_message = []


def expire_user_tasks():
    user_balance_count = cloud.balances.count_documents({"uid": {"$exists": True}})
    unexpire_balance_count = cloud.balances.count_documents({'planExpireTime': {'$exists': True, '$gt': arrow.now().shift(days=-1).datetime}})
    will_expire_uids = cloud.balances.distinct('uid', {
        'planExpireTime': {'$exists': True, '$gt': arrow.now().datetime, "$lt": arrow.now().shift(days=7).datetime}})
    will_expire_usernames = cloud.users.distinct("username", {"_id": {"$in": will_expire_uids}})
    expire_uids = cloud.balances.distinct('uid', {'planExpireTime': {'$exists': True, '$lt': arrow.now().shift(days=-1).datetime}})
    asset_task_count = cloud.asset_tasks.count_documents({
        'triggerType': 'schedule', 'enableMonitor': True, 'uid': {'$in': expire_uids}
    })
    ats = cloud.asset_tasks.find({'triggerType': 'schedule', 'enableMonitor': True, 'uid': {'$in': expire_uids}}, {
        'createTime': 1,
        'vulSettings.taskId': 1,
        'hostVulSettings.taskId': 1,
        'sslSettings.taskId': 1,
        'securityEventSettings.taskId': 1,
        'contentSettings.taskId': 1,
        'assetSettings.taskId': 1,
        'httpSettings.taskId': 1,
        'pingSettings.taskId': 1,
        'ipv6Settings.taskId': 1,
    })
    task_ids = []
    for at in ats:
        for k, v in at.items():
            if 'Settings' not in k:
                continue
            tid = v.get('taskId')
            if tid:
                task_ids.append(tid)
    query = {
        'task_id': {'$in': task_ids},
        'task_config.is_disabled': False,
        'task_config.is_periodic': True,
        '$or': [
            {'addition.notificationTarget': {'$exists': 0}},
            {'addition.notificationTarget': 1},
        ]
    }
    scan_tasks = wass.scan_tasks.distinct('task_id', query)
    uid = wass.scan_tasks.distinct('addition.uid', query)
    user_info = {}
    for user in cloud.users.find({'_id': {'$in': [ObjectId(_) for _ in uid]}}):
        balance = cloud.balances.find_one({'_id': user['balanceId']})
        user_info[user['username']] = balance.get('planExpireTime', '') if balance else ''

    info = [f'{k} {str(v)}' for k, v in user_info.items()]
    notify_message.append(
        f"注册用户数:{user_balance_count} 未过期:{unexpire_balance_count} 已过期:{len(expire_uids)} 七日内即将过期用户:[{','.join(will_expire_usernames)}],"
        f"前台仍开启周期任务{asset_task_count}个, 后台遗留任务{len(scan_tasks)}个, 遗留用户数{len(user_info)}, 用户信息{info}"
    )
    logger.info(notify_message[-1])


def clean_wass_tasks(remove=False):
    expire_uids = cloud.balances.distinct('uid', {'planExpireTime': {'$exists': True, '$lt': arrow.now().shift(days=-3).datetime}})
    user_expire_map = {str(_['_id']): True if _['_id'] in expire_uids else False for _ in cloud.users.find()}
    tasks = wass.scan_tasks.find({
        'task_config.is_disabled': False,
        'task_config.is_periodic': True,
        '$or': [
            {'addition.notificationTarget': {'$exists': 0}},
            {'addition.notificationTarget': 1},
        ]
    }, {'addition': 1, 'task_id': 1, 'target_url': 1})
    stop_tasks = []
    for _ in tasks:
        if not _['addition']:
            continue
        uid = _['addition']['uid']
        try:
            r = user_expire_map[uid]
            if r:
                stop_tasks.append(_['task_id'])
        except KeyError:
            logger.info(f'Not found {str(uid)} {_["task_id"]} {_["target_url"]} {_["addition"].get("taskType", "")}')
    removed = 0
    # 每次删除5个task
    count = 5
    if remove:
        for index in range(math.ceil(len(stop_tasks) / count)):
            task_ids = stop_tasks[index*count:(index+1)*count]
            request_args = {
                'taskIds': task_ids
            }
            resp = monitor.post('/api/v2/jobs/tasks/scanner/remove_scanner_tasks', request_args)
            if resp and resp.status_code == 200:
                logger.info(f'Remove tasks {task_ids}')
                removed += len(task_ids)
                time.sleep(5)
            else:
                logger.info(f'Remove error {resp}')
    notify_message.append(f'需要删除已过期用户任务{len(stop_tasks)}, 实际删除任务{removed}')
    logger.info(notify_message[-1])


def check_task_enable():
    scan_task_ids = wass.scan_tasks.distinct('task_id', {
        'task_config.is_disabled': False,
        'task_config.is_periodic': True,
        'addition': {'$exists': 1, '$ne': {}},
        'addition.taskSettings.enableSitePortraitTriggerMonitor': {'$ne': True},
        '$or': [
            {'addition.notificationTarget': {'$exists': 0}},
            {'addition.notificationTarget': 1},
        ]
    })
    task_type_count = ', '.join([f'{_["_id"]}:{_["sum"]}' for _ in wass.scan_tasks.aggregate([
        {'$match': {
            'task_config.is_disabled': False,
            'task_config.is_periodic': True,
            'addition': {'$exists': 1, '$ne': {}},
            'addition.taskSettings.enableSitePortraitTriggerMonitor': {'$ne': True},
            '$or': [
                {'addition.notificationTarget': {'$exists': 0}},
                {'addition.notificationTarget': 1},
            ]
        }},
        {'$group': {'_id': '$addition.taskType', 'sum': {'$sum': 1}}}
    ])])
    notify_message.append(f'引擎开启周期任务数{len(scan_task_ids)}\n{task_type_count}')
    logger.info(notify_message[-1])
    ats = cloud.asset_tasks.find({'triggerType': 'schedule', 'enableMonitor': True}, {
        'uid': 1,
        '_id': 1,
        'hostVulSettings': 1,
        'vulSettings': 1,
        'sslSettings': 1,
        'securityEventSettings': 1,
        'contentSettings': 1,
        'assetSettings': 1,
        'httpSettings': 1,
        'pingSettings': 1,
        'ipv6Settings': 1,
    })
    enable_task_map = dict()
    enable_tasks = set()
    for at in ats:
        _id = at.get("_id")
        for k, v in at.items():
            if 'Settings' not in k:
                continue
            tid = v.get('taskId')
            enable = v.get('enable')
            if tid and enable:
                # 开启网站画像的任务，在引擎可能为周期或单次任务
                if v.get("enableSitePortraitTriggerMonitor"):
                    continue
                enable_tasks.add(tid)
                enable_task_map.update({tid: _id})
    notify_message.append(f'前台开启周期任务数{len(enable_tasks)}')
    logger.info(notify_message[-1])
    ret = [_ for _ in scan_task_ids if _ not in enable_tasks]
    notify_message.append(f'引擎开启周期任务，前台任务未开启或无任务{len(ret)}')
    logger.info(notify_message[-1])

    # 处理前台无job任务
    for _ in wass.scan_tasks.find({'task_id': {'$in': ret}}):
        try:
            job = cloud.jobs.find_one({'_id': ObjectId(_['addition']['jobId'])})
            if job:
                logger.info(f"{job['_id']} {job['targetUrl']} {_['task_id']} {_['addition']['taskType']}")
            else:
                logger.info(f"{_['addition']['jobId']} {_['task_id']} {_['addition']['taskType']}")
    #             task_id = _['task_id']
    #             request_args = {'taskIds': [task_id]}
    #             resp = monitor.post('/api/v2/jobs/tasks/scanner/remove_scanner_tasks', request_args)
    #             if resp and resp.status_code == 200:
    #                 logger.info(f'Remove tasks {task_id}')
    #             else:
    #                 logger.info(f'Remove error {resp}')
        except Exception as e:
            logger.info(e)
            logger.info(_['task_id'])

    # 处理asset缺陷未开启enableMonitor
    # asset = cloud.asset_tasks.count_documents({'assetSettings.taskId': {'$in': ret}})
    # print(f'处理变更监测{asset}个')
    # for _ in ret:
    #     a = cloud.asset_tasks.find_one({'assetSettings.taskId': _})
    #     if a and not a['enableMonitor']:
    #         cloud.asset_tasks.update_one({'assetSettings.taskId': _}, {'$set': {'enableMonitor': True}})
    #         print(_)

    no_wass = [_ for _ in enable_tasks if _ not in scan_task_ids]
    notify_message.append(f'前台开启周期任务，引擎任务未开启或无任务{len(no_wass)}')
    logger.info(notify_message[-1])
    logger.info([(no_tid, enable_task_map.get(no_tid, "")) for no_tid in no_wass])
    # for _ in wass.scan_tasks.find({'task_id': {'$in': no_wass}}, {'addition.taskType': 1, 'task_id': 1}):
    #     print(_['addition']['taskType'], _['task_id'])
    # user_map = {_['_id']: _['username'] for _ in cloud.users.find()}
    # for task_type in ['vul', 'ssl', 'content', 'securityEvent', 'http', 'ping']:
    #     for _ in cloud.asset_tasks.find({f'{task_type}Settings.taskId': {'$in': no_wass}}):
    #         print(_['_id'], _['jobId'], _[f'{task_type}Settings']['taskId'], task_type)
    #         print(_['targetUrl'], _['uid'], user_map.get(_['uid'], ''))


def find_repeat_task():
    # TODO
    key = {
        'vulSettings.taskId': 1,
        'hostVulSettings.taskId': 1,
        'sslSettings.taskId': 1,
        'securityEventSettings.taskId': 1,
        'contentSettings.taskId': 1,
        'assetSettings.taskId': 1,
        'httpSettings.taskId': 1,
        'pingSettings.taskId': 1,
        'ipv6Settings.taskId': 1,
    }
    for k, v in key.items():
        print(k)
        ret = cloud.asset_tasks.aggregate([
            {'$match': {f'{k}': {'$exists': 1, '$ne': ''}, 'triggerType': 'schedule', 'enableMonitor': True}},
            {'$group': {'_id': f'${k}', 'sum': {'$sum': 1}}},
            {'$match': {'sum': {'$gt': 1}}}
        ])
        for _ in ret:
            url = cloud.asset_tasks.distinct('jobId', {f'{k}': _['_id'], 'enableMonitor': True})
            if len(url) > 1:
                print(url, _['_id'])


if __name__ == "__main__":
    expire_user_tasks()
    clean_wass_tasks(remove=False)
    check_task_enable()
    # logger.info('\n'.join(notify_message))
    notify('\n'.join(notify_message))
    # find_repeat_task()
